import argparse
import os
import cv2
import sys
import xml.etree.ElementTree as ET

#def process_image(img,min_side):

path  = 'crane_new'


if __name__ == '__main__':
    imgFilePath = os.path.join(path, 'JPEGImages')
    dstImgFilePath = os.path.join(path, 'JPEGImages_com')

    xmlFilePath = os.path.join(path, 'Annotations')
    dstXmlFilePath = os.path.join(path, 'Annotations_com')

    #检查输入目录
    if not os.path.exists(dstImgFilePath):
        os.mkdir(dstImgFilePath)   
    if not os.path.exists(dstXmlFilePath):
        os.mkdir(dstXmlFilePath)

    #遍历xml文件
    xmlFiles = os.listdir(xmlFilePath)
    for fileName in xmlFiles:
        imgFileName = os.path.splitext(fileName)[0]+".jpg"
        imgFile = os.path.join(imgFilePath,imgFileName)
        dstImgFile = os.path.join(dstImgFilePath,imgFileName)
        xmlfile = os.path.join(xmlFilePath,fileName)
        dstXmlFile = os.path.join(dstXmlFilePath,fileName)

        #压缩处理
        img = cv2.imread(imgFile)
        size = img.shape
        h,w = size[0],size[1]

        #短边小于832不压缩
        if min(h,w)<832:
            continue

        #短边缩放为目标尺寸
        scale = min(w,h)/float(416)
        new_w,new_h = int(w/scale),int(h/scale)
        resize_img = cv2.resize(img,(new_w,new_h))
        cv2.imwrite(dstImgFile,resize_img)
        print(dstImgFile)
        
        #调整目标框
        tree = ET.ElementTree
        tree = ET.parse(xmlfile)
        root = tree.getroot()
        for obj in root.findall('object'):
            bndbox  = obj.find('bndbox')
            left    = int(bndbox.find('xmin').text)
            top     = int(bndbox.find('ymin').text)
            right   = int(bndbox.find('xmax').text)
            bottom  = int(bndbox.find('ymax').text)
            left = int(left/scale)
            top = int(top/scale)
            right = int(right/scale)
            bottom = int(bottom/scale)
            bndbox.find('xmin').text = str(left)
            bndbox.find('ymin').text = str(top)
            bndbox.find('xmax').text = str(right)
            bndbox.find('ymax').text = str(bottom)
        tree.write(dstXmlFile,encoding="utf-8", xml_declaration=False)       
        print(dstXmlFile)     